import os
import sys
import random
import pprint

import pygame
from pygame.locals import *

class GeneticCode:
  """Genetic Code"""
        
  def __init__(self, translationTable = None, startCodon = '00', stopCodon = '10', lenCodon = 2):
    if translationTable is None:
        self.translationTable = {'00': 'H',
                                 '01': '-',
                                 '10': '*',
                                 '11': 'V'}
    else:
      self.translationTable = translationTable
      
    self.startCodon = startCodon
    self.stopCodon = stopCodon
    self.lenCodon = lenCodon
      
  def isStartCodon(self, codon):
    return codon == self.startCodon
      
  def isStopCodon(self, codon):
    return codon == self.stopCodon
      
  def translateCodon(self, codon):
    if len(codon) == self.lenCodon:
      return self.translationTable[codon]
    else:
      return None

  def translateSeq(self, seq):
    translating = False
    proteina = []
    for codon in seq:
        if translating:
            if self.isStopCodon(codon):
                return proteina
            else:
                amino = self.translateCodon(codon)
                if not amino is None:
                  proteina.append(amino)
        else:
            if self.isStartCodon(codon):
                proteina.append(self.translateCodon(codon))
                translating = True
    
    if len(proteina) == 0:
      proteina = None
      
    return proteina

class ProteinBag:
  """Bolsa de proteinas"""
  
  def __init__(self):
    self.bag = {}
    self.volume = 0
  
  def append(self, proteina):
    self.appendN(proteina, 1)
  
  def appendN(self, proteina, n):
    self.volume += n * len(proteina)
    if self.bag.has_key(proteina):
      self.bag[proteina] = self.bag[proteina] + n
    else:
      self.bag[proteina] = n
  
  def substractN(self, proteina, n):
    if self.bag.has_key(proteina):
      self.bag[proteina] = self.bag[proteina] - n
      self.volume -= len(proteina)
    
#    if self.bag[proteina] < 0:
#      del self.bag[proteina]

  def substract(self, proteina):
    self.substractN(proteina, 1)

  def getAmountOfProtein(self, proteina):
    if self.bag.has_key(proteina):
      return self.bag[proteina]
    else:
      return 0  
  
  def getVolume(self):
    """ devuelve el volumen de la bolsa"""
    return self.volume
  
  def info(self):
    txt_bag = "[" + ", ".join(map(lambda x: str(self.bag[x]) + ' ' + str(x), self.bag)) + "]"
    return txt_bag

    
class Membrane:
  """Membrana celular"""
  
  def __init__(self, cell):
    self.cell = cell
    self.freeProteins = ProteinBag()
    self.joinProteins = []
  
  def __str__(self):
    return "Membrana %d" % self.cell.id
  
  def doJoin(self, protein, membrane):
    self.joinProteins.append((protein, membrane))
    self.freeProteins.substract(protein)
  
  def interact(self, other_membrane):
    num_interacts = 0
    # optimization, loop over shorter dict ;-)
    if len(other_membrane.freeProteins.bag) < len(self.freeProteins.bag):
      other_membrane.interact(self)
    else:
      for protein, n in self.freeProteins.bag.iteritems():
        if n > 0 and other_membrane.freeProteins.getAmountOfProtein(protein) > 0:
          num_interacts += 1
          self.doJoin(protein, other_membrane)
          other_membrane.doJoin(protein, self)
          
    return num_interacts

        
class Cell:
  """Cell class"""
  id = 0
  
  def __init__(self, citizen, cg, adn):
    Cell.id = Cell.id + 1
    self.id = Cell.id
    
    self.citizen = citizen
    
    self.cg = cg    
    
    self.adn = adn
    self.adnCopy = ''
    
    self.proteinBag = ProteinBag()
    
    self.proteinBagCopy = None
    self.freeProteinsCopy = None
    self.joinProteinsCopy = None
    
    self.membrana = Membrane(self)
    
    self.fase = 'G1'

  def __str__(self):
    return "Cell %d" % self.id
  
  def info(self):
    txt_cell = str(self)
    txt_fase = "  Fase: %s" % self.fase
    txt_volume = "  Volumen: %d" % self.proteinBag.getVolume()
    txt_membrana = "  Membrana:"
    txt_free = "    free: [" + ", ".join(map(lambda x: str(self.membrana.freeProteins.bag[x]) + ' ' + str(x), self.membrana.freeProteins.bag)) + "]"
    txt_join = "    join: [" + ", ".join(map(lambda x: '(' + str(x[0]) + ', ' + str(x[1]) + ')', self.membrana.joinProteins)) + "]"
    txt_proteinas = "  Proteinas: " + self.proteinBag.info()
    return [txt_cell, txt_fase, txt_volume, txt_membrana, txt_free, txt_join, txt_proteinas]
    
  def cicle(self):
    if self.fase == 'G1':
      self.cicleG1()
      
    elif self.fase == 'S':
      self.cicleS()
      
    elif self.fase == 'G2':
      self.cicleG2()
      
    elif self.fase == 'M':
      self.cicleM()
      
  def cicleM(self):
    """Fase M division celular"""
    # option 1: fuerza una primera interaccion entre las celulas recien nacidas
    # option 2: dejar la celula recien nacida como citizen independiente
    #           pero en la misma posicion para que se suja por si mimsa la
    #           posible interaccion. Esto tiene la ventaja de que la recien
    #           nacida podria unirse a cualquier otra celula del citizen
    # de momento usamos la opcion 2
    citizen = Citizen(self.citizen.world)
    citizen.pos = self.citizen.pos

    cell = Cell(citizen, self.cg, self.adnCopy)
    
    for proteina, num in self.proteinBagCopy.bag.iteritems():
      cell.proteinBag.appendN(proteina, num)
    self.proteinBagCopy = None

    for proteina, num in self.freeProteinsCopy.bag.iteritems():
      cell.membrana.freeProteins.appendN(proteina, num)   
    self.freeProteinsCopy = None
    
    # TODO IMPORTANT: fix joins!!!
    # la celula opuesta no sabe que ahora esta enganchada a otra celula
    # diferente.
    cell.membrana.joinProteins = self.joinProteinsCopy
    self.joinProteinsCopy = None
    
    citizen.cells = [cell]
    citizen.world.appendCitizen(citizen)      
    
    self.fase = 'G1'
#    print "%d -> changed to fase G1!" % self.id
    
  def cicleG2(self):
    """Fase G2 reparto de contenido"""        
    #
    # primero se reparten las proteinas del interior
    # 50% de posibilidades de que cada proteina recaiga
    # en una u otra copia de la celula
    #
    self.proteinBagCopy = ProteinBag()
    for p,v in self.proteinBag.bag.iteritems():
      for i in range(v):
        if random.random() < 0.5:
          self.proteinBagCopy.append(p)
          
    for p,v in self.proteinBagCopy.bag.iteritems():      
      self.proteinBag.substractN(p, v)
      
    #
    # despues repartimos las proteinas libres de la membrana
    #
    self.freeProteinsCopy = ProteinBag()
    for p,v in self.membrana.freeProteins.bag.iteritems():
      for i in range(v):
        if random.random() < 0.5:
          self.freeProteinsCopy.append(p)
          
    for p,v in self.freeProteinsCopy.bag.iteritems():    
      self.membrana.freeProteins.substractN(p, v)
    
    #
    # finalmente repartimos las proteins unidas de la membrana
    # 
    self.joinProteinsCopy = []
    for join in self.membrana.joinProteins:
      if random.random() < 0.5:
        self.joinProteinsCopy.append(join)
        self.membrana.joinProteins.remove(join)
      
#    print "%d -> changed to fase M!" % self.id
    self.fase = 'M'
    
  def cicleS(self):
    """Fase S duplicacion del adn"""
    pos = len(self.adnCopy)
    n = self.adn[pos:pos+1]
    self.adnCopy += n
    
    if len(self.adnCopy) >= len(self.adn):
#      print "%d -> changed to fase G2!" % self.id    
      self.fase = 'G2'
      
    return
    
  def cicleG1(self):
    """Fase G1 engorde"""
    if self.proteinBag.getVolume() > 10000:
#      print "%d -> changed to fase S!" % self.id
      self.fase = 'S'
      return
    
    # select random start point in adn
    start = random.randint(0, len(self.adn) - 1)
    
    # build sequence to translate
    if start == 0:
        plainSeq = self.adn
    else:
        plainSeq = self.adn[start:] + self.adn[:start-1]
        
    # break seq in codons
    seq = list(self.codons(plainSeq))
    
    # translate into codons
    aminoacidos = self.cg.translateSeq(seq)
    
    # and store it
    if not aminoacidos is None:
      proteina = "".join(aminoacidos) # plegar?
      if proteina[-1] == "-":
        self.membrana.freeProteins.append(proteina)
      else:
        self.proteinBag.append(proteina)
          
  def codons(self,seq):
    for i in xrange(0, len(seq), self.cg.lenCodon):
      yield seq[i:i + self.cg.lenCodon]

class Citizen:
  id = 0
  
  def __init__(self, world):
    Citizen.id = Citizen.id + 1
    self.id = Citizen.id
    self.world = world
    self.cells = []
    self.pos = 0, 0
  
  def __str__(self):
    return "Citizen %d" % self.id
    
  def info(self):
    info = []
    info.append(str(self))
    
    info_cells = []
    for cell in self.cells:
      info_cell = map(lambda x: "  " + x, cell.info())
      info.extend(info_cell)
      
    return info
    
  def cicle(self):
    for cell in self.cells:
      cell.cicle()
      
  def selectRandomCell(self):
    i = random.randint(0, len(self.cells) - 1)
    return self.cells[i]
    
  def move(self):
    x_offset = float(1 - random.randint(0,2)) / len(self.cells)
    y_offset = float(1 - random.randint(0,2)) / len(self.cells)
    
    x, y = self.pos    
    x += x_offset / len(self.cells)
    y += y_offset / len(self.cells)
    
    if x < 0:
      x = self.world.maxX
    if x > self.world.maxX:
      x = 0
    if y < 0:
      y = self.world.maxY
    if y > self.world.maxY:
      y = 0
    
    self.pos = x, y
    
  def getJoints(self):
    joints = {}
    for cell in self.cells:
      cjoints = cell.getJoints();
      if not cjoints is None:
        for k, v in cjoints.iteritems():
          if joints.has_key(k):          
            n = joints[k]
            joints[k] = max(n, v) - min(n, v)
          else:
            joints[k] = cjoints[k]
        
    return joints
          
class World:
  
  def __init__(self, num_citizens = 100, len_adn = 100):
    self.cg = GeneticCode()    
    self.maxX = 400
    self.maxY = 400
    self.poss = {}
    self.cicleNum = 0
    self.newCitizens = []
    
    # generate initial unicellular aleat population
    self.population = []
    for i in range(num_citizens):
      adn = "".join(['0' if random.random() < 0.5 else '1' for x in range(len_adn)])
      
      citizen = Citizen(self)
      citizen.pos = random.randint(0, self.maxX), random.randint(0, self.maxY)
      
      cell = Cell(citizen, self.cg, adn)
      for j in range(100):
        cell.cicle()
        
      cells = [cell]

      citizen.cells = cells
      
      self.population.append(citizen)
  
  def appendCitizen(self, citizen):
    self.newCitizens.append(citizen)
    
  def cicle(self):
    self.cicleNum += 1
    self.poss = {}
    for citizen in self.population:
      citizen.cicle()
      if citizen.pos in self.poss:
        (self.poss[citizen.pos]).append(citizen)
      else:
        self.poss[citizen.pos] = [citizen]
        
    for citizen in self.newCitizens:
      self.population.append(citizen)
      
    self.newCitizens = []
     
    #
    # Interactions
    #
    # When 2 citizens fall in the same position they interact,
    # all of them interact 2 by 2.
    #
    to_remove = []
    to_add = []
    for pos, vecinos in self.poss.iteritems():
      if len(vecinos) >= 2:
        pairs = [(x,y) for x in vecinos for y in vecinos if x <> y]       
        for pair in pairs:                    
#          print "-> interact..." + str(pair[0]) + " - " + str(pair[1])
          new_citizen = self.interact(pair[0], pair[1])
          if not new_citizen is None:
            print "-> crossing..." + str(pair[0]) + " - " + str(pair[1])
            # print ", ".join(map(str, self.population))
            
            to_remove.append(pair[0])
            to_remove.append(pair[1])
            to_add.append(new_citizen)
            
    for c in to_remove:
      if c in self.population:
        self.population.remove(c)

    for c in to_add:
      if not c in self.population:
        self.population.append(c)

    # move citizens
    for citizen in self.population:
      citizen.move()    
  
  def interact(self, citizen1, citizen2):
    cell1 = citizen1.selectRandomCell()
    cell2 = citizen2.selectRandomCell()
    
    # solo se permite interaccion en la fase de la celula G1
    if cell1.fase != 'G1' or cell2.fase != 'G1':
#      print "   do not interact: selected cells not in G1."
      return None
      
    # y no en celulas que ya esten unidas entre si
    if cell1 in citizen2.cells or cell2 in citizen1.cells:
#      print "   do not interact: selected cells already joined."
      return None
      
    num = cell1.membrana.interact(cell2.membrana)
#    print "   do interact: on %d proteins!" % num
    
    if num > 0:
      citizen1.cells.extend(citizen2.cells)
      return citizen1
    else:
      return None
    
  def old_interact(self, citizen1, citizen2):
    js1 = citizen1.getJoints()
    js2 = citizen2.getJoints()
    joint_value = 0
    
    if len(js1) == 0 or len(js2) == 0:
      return None

    for k, v in js1.iteritems():
      if js2.has_key(k):
        joint_value += min(v, js2[k]) * len(k)
        
    # compare joint_value with arbitrary fixed value
    if joint_value > 10:
      pp = pprint.PrettyPrinter(indent=4)
    
      print "Interacting...%d" % joint_value
      print "--------------"
      pp.pprint((js1, js2))
      # pp.pprint(js2)
      print "--------------"
      citizen1.cells.extend(citizen2.cells)
      return citizen1
    else:
      return None    
    
  def printPopulation(self):
#    pp = pprint.PrettyPrinter(indent=4)
#    pp.pprint (list(map(str, self.population)))
#    print "\n".join(map(info, self.population))
    for citizen in self.population:
      print "\n".join(map(str, citizen.info()))
    
    # for citizen in self.population:    
      # pp.pprint([citizen.id, citizen.pos])
      # pp.pprint(citizen.cells)
        
    # pp.pprint(self.poss)

def main(argv=None):
  pygame.init() 
  window = pygame.display.set_mode((400, 400)) 

  background = pygame.Surface(window.get_size())
  background = background.convert()
  background.fill((5, 5, 5))
  centerx = background.get_width()/2
  
  w = World(100,100)
  for i in range(2500):
    pygame.display.flip()
    
    w.cicle()


    window.blit(background, (0, 0))
    
    if pygame.font:
        font = pygame.font.Font(None, 20)
        infocicle = "Ciclo: %d" % w.cicleNum
        textcicle = font.render(infocicle, 1, (255, 255, 10))
        
        infocitizens = " Citizens: %d" % len(w.population)
        textcitizens = font.render(infocitizens, 1, (255, 255, 10))
        
        window.blit(textcicle, (0,0))
        window.blit(textcitizens, (100,0))
        
    for citizen in w.population:
      pos = citizen.pos
      for cell in citizen.cells:
        size = cell.proteinBag.getVolume() / 1000
        pygame.draw.circle(window, (128, 128, 255), (int(pos[0]), int(pos[1])), size)
        pos = (pos[0] + size, pos[1])
      
    
  w.printPopulation()
#  os.system("PAUSE")  

if __name__ == "__main__":
    sys.exit(main())
